package edu.cornell.cs4740.qa.ir;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.Field.Index;
import org.apache.lucene.document.Field.Store;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.IndexDeletionPolicy;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriter.MaxFieldLength;
import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.FuzzyQuery;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.store.LockObtainFailedException;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.util.Version;

public class DocumentQuery {

	Directory indexDirectory;
	Analyzer analyzer;
	IndexSearcher iSearcher;

	public DocumentQuery(edu.cornell.cs4740.qa.Document document, boolean useRamBasedIndex)
			throws CorruptIndexException, LockObtainFailedException, IOException {
		if (useRamBasedIndex) {
			indexDirectory = new RAMDirectory();
		} else {
			try {
				indexDirectory = FSDirectory.open(new File("./index"));
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}

		analyzer = new StandardAnalyzer(Version.LUCENE_36);

		IndexDeletionPolicy idp = new KeepOnlyLastCommitDeletionPolicy();

		IndexWriter iWriter = new IndexWriter(indexDirectory, analyzer, idp, MaxFieldLength.UNLIMITED);

		List<String> sentences = document.getSentencesStrings();

		for (String sentence : sentences) {
			// System.out.println(sentence);
			org.apache.lucene.document.Document sentenceDoc = new Document();
			Field sentenceField = new Field("SENTENCE", sentence, Store.YES, Index.ANALYZED);
			Field docIdField = new Field("DOC_ID", document.getDocId(), Store.YES, Index.NO);
			sentenceDoc.add(sentenceField);
			sentenceDoc.add(docIdField);
			iWriter.addDocument(sentenceDoc);
		}

		iWriter.optimize();
		iWriter.close();

		IndexReader iReader = IndexReader.open(indexDirectory);
		iSearcher = new IndexSearcher(iReader);
	}

	public DocumentQuery(List<edu.cornell.cs4740.qa.Document> documents, boolean useRamBasedIndex)
			throws CorruptIndexException, LockObtainFailedException, IOException {
		if (useRamBasedIndex) {
			indexDirectory = new RAMDirectory();
		} else {
			try {
				indexDirectory = FSDirectory.open(new File("./index"));
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}

		analyzer = new StandardAnalyzer(Version.LUCENE_36);

		IndexDeletionPolicy idp = new KeepOnlyLastCommitDeletionPolicy();

		IndexWriter iWriter = new IndexWriter(indexDirectory, analyzer, idp, MaxFieldLength.UNLIMITED);

		List<String> sentences = new ArrayList<String>();

		for (edu.cornell.cs4740.qa.Document document : documents) {
			sentences.addAll(document.getSentencesStrings());
		}

		for (String sentence : sentences) {
			// System.out.println(sentence);
			org.apache.lucene.document.Document sentenceDoc = new Document();
			Field sentenceField = new Field("SENTENCE", sentence, Store.YES, Index.ANALYZED);
			// Field docIdField = new Field("DOC_ID", document.getDocId(),
			// Store.YES, Index.NO);
			sentenceDoc.add(sentenceField);
			// sentenceDoc.add(docIdField);
			iWriter.addDocument(sentenceDoc);
		}

		iWriter.optimize();
		iWriter.close();

		IndexReader iReader = IndexReader.open(indexDirectory);
		iSearcher = new IndexSearcher(iReader);
	}

	public List<String> getMostRelevantSentences(List<String> searchTerms, int numSentencesToReturn)
			throws CorruptIndexException, IOException {
		List<String> relevantSentences = new ArrayList<String>();

		// PhraseQuery query = new PhraseQuery();

		BooleanQuery query = new BooleanQuery();

		for (String t : searchTerms) {
			query.add(new FuzzyQuery(new Term("SENTENCE", t)), Occur.SHOULD);
		}

		// query.setSlop(SEARCH_SLOP);

		TopDocs topSentences = iSearcher.search(query, numSentencesToReturn);

		ScoreDoc[] scoreDocs = topSentences.scoreDocs;

		for (int x = 0; x < scoreDocs.length; x++) {
			ScoreDoc sd = scoreDocs[x];
			int docNum = sd.doc;

			Document relevantDoc = iSearcher.doc(docNum);

			relevantSentences.add(relevantDoc.getField("SENTENCE").stringValue());
		}

		return relevantSentences;
	}
}
